#!/user/bin/env python
# -*- coding:utf-8 -*-
# 作者：洛月
# 创建：2021-03-11 17:17
# 更新：2021-03-11 17:17
# 功能：
import os
import cv2
import numpy
import time
import sqlite3
import imghdr
from PIL import Image


class ColorDescriptor:
    __slot__ = ["bins"]
    bins = (8, 12, 3)

    # 得到图片的色彩直方图，mask为图像处理区域的掩模
    def getHistogram(self, image, mask, isCenter):
        # 利用OpenCV中的calcHist得到图片的直方图
        imageHistogram = cv2.calcHist([image], [0, 1, 2], mask, self.bins, [0, 180, 0, 256, 0, 256])
        # 标准化(归一化)直方图normalize
        imageHistogram = cv2.normalize(imageHistogram, imageHistogram).flatten()
        # isCenter判断是否为中间点，对色彩特征向量进行加权处理
        if isCenter:
            weight = 5.0  # 权重记为0.5
            for index in range(len(imageHistogram)):
                imageHistogram[index] *= weight
        return imageHistogram

    # 将图像从BGR色彩空间转换为HSV色彩空间
    def describe(self, image):
        image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
        features = []
        # 获取图片的中心点和图片的大小
        height, width = image.shape[0], image.shape[1]
        centerX, centerY = int(width * 0.5), int(height * 0.5)
        # initialize mask dimension
        # 生成左上、右上、左下、右下、中心部分的掩模。
        # 中心部分掩模的形状为椭圆形。这样能够有效区分中心部分和边缘部分，从而在getHistogram()方法中对不同部位的色彩特征做加权处理。
        segments = [(0, centerX, 0, centerY), (0, centerX, centerY, height), (centerX, width, 0, centerY),
                    (centerX, width, centerY, height)]
        # 初始化中心部分
        axesX, axesY = int(width * 0.75) / 2, int(height * 0.75) / 2
        ellipseMask = numpy.zeros([height, width], dtype="uint8")
        cv2.ellipse(ellipseMask, (int(centerX), int(centerY)), (int(axesX), int(axesY)), 0, 0, 360, 255, -1)
        # cv2.ellipse(ellipMask, (int(cX), int(cY)), (int(axesX), int(axesY)), 0, 0, 360, 255, -1)
        # 初始化边缘部分
        for startX, endX, startY, endY in segments:
            cornerMask = numpy.zeros([height, width], dtype="uint8")
            cv2.rectangle(cornerMask, (startX, startY), (endX, endY), 255, -1)
            cornerMask = cv2.subtract(cornerMask, ellipseMask)
            # 得到边缘部分的直方图
            imageHistogram = self.getHistogram(image, cornerMask, False)
            features.append(imageHistogram)
        # 得到中心部分的椭圆直方图
        imageHistogram = self.getHistogram(image, ellipseMask, True)
        features.append(imageHistogram)
        # 得到最终的特征值
        return features


# 将图片进行归一化处理，返回HSV色彩空间矩阵
class StructureDescriptor:
    __slot__ = ["dimension"]
    dimension = (16, 16)

    def describe(self, image):
        image = cv2.resize(image, self.dimension, interpolation=cv2.INTER_CUBIC)
        # 将图片转化为BGR图片转化为HSV格式
        image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
        # print(image)
        return image


imgFormat = ['png', 'jpeg']


def uploadImage(file):
    if imghdr.what(file) in imgFormat:
        if not os.path.exists('images'):
            os.mkdir('images')
        try:
            C = ColorDescriptor()
            S = StructureDescriptor()
            imgId = int('1' + str(time.time()).replace('.', '')[-6:])
            Image.open(file).save(f'images/{imgId}.png')
            img = cv2.imread(f'images/{imgId}.png')
            db = sqlite3.connect('sqlite3.db')
            cur = db.cursor()
            cur.execute("insert into color(id,v1,v2,v3,v4,v5) values(?,?,?,?,?,?)",
                        tuple([imgId] + [str(i).replace('\n', '') for i in C.describe(img)]))
            cur.execute(
                "insert into structure(id,v1,v2,v3,v4,v5,v6,v7,v8,v9,v10,v11,v12,v13,v14,v15,v16) values(?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)",
                tuple([imgId] + [str(i).replace('\n', '') for i in S.describe(img)]))
            db.commit()
            return imgId
        except Exception as e:
            return False
    else:
        return False


if __name__ == '__main__':
    for i in os.listdir(r'F:\Python\cv\get'):
        print(uploadImage(os.path.join(r'F:\Python\cv\get', i)))
        time.sleep(1)
